from openai import OpenAI


def moonshot_one(data):
    """
    moon_shot v1 8k 接口，接收一个data，是一个dict对象，详细如下：
    data: {
        "message": [] # 列表，
        "config": {
            "temperature": 0-2, default 1，
            "max_token": 0-1024, default 512
        }
    }
    """

    client = OpenAI(
        api_key="sk-x8A8bpJALwyLlE1J8noYBfa9PaNeRXs2fi0xZ5Y5M5VDhsfy",
        base_url="https://api.moonshot.cn/v1",
    )
    completion = client.chat.completions.create(
        model="moonshot-v1-8k",

        # default 512, between 0-1024
        max_tokens=data["config"]["max_token"] if (data["config"]["max_token"] >= 0) and (
                data["config"]["max_token"] <= 1024) else 512,

        # default 1, between 0 and 2
        temperature=data["config"]["temperature"] if (data["config"]["temperature"] >= 0) and (
                data["config"]["temperature"] <= 1) else 0.5,

        # 如果有system,则在列表索引为0处
        messages=data["message"]
    )

    # print(completion.choices[0].message.content)
    return completion.choices[0].message.content


def moonshot_one_stream(data):
    client = OpenAI(
        api_key="sk-x8A8bpJALwyLlE1J8noYBfa9PaNeRXs2fi0xZ5Y5M5VDhsfy",
        base_url="https://api.moonshot.cn/v1",
    )
    completion = client.chat.completions.create(
        model="moonshot-v1-8k",

        # default 512, between 0-1024
        max_tokens=data["config"]["max_token"] if (data["config"]["max_token"] >= 0) and (
                data["config"]["max_token"] <= 1024) else 512,

        # default 1, between 0 and 2
        temperature=data["config"]["temperature"] if (data["config"]["temperature"] >= 0) and (
                data["config"]["temperature"] <= 1) else 0.5,

        # 如果有system,则在列表索引为0处
        messages=data["message"],
        stream=True  # 流式返回
    )

    for idx, chunk in enumerate(completion):
        chunk_message = chunk.choices[0].delta
        if not chunk_message.content:
            continue
        print(chunk_message.content)


if __name__ == "__main__":
    import os

    # store openai-api-key to internal storage
    # os.environ["OPENAI_API_KEY"] = "sk-proj-wOXa1EvcrzPxGvVwjbp3T3BlbkFJSZNHhSdxIjtB02yFquNW"

    # solve proxy problem
    os.environ["http_proxy"] = "http://localhost:7890"
    os.environ["https_proxy"] = "http://localhost:7890"

    data1 = {
        "message": [
            {"role": "system", "content": "你是一个有用的助手。"},
            {"role": "user", "content": "你好啊"}
        ],
        "config": {
            "temperature": 1,
            "max_token": 512
        }
    }

    moonshot_one_stream(data1)
